# -*- coding: utf-8 -*-
from typing import Callable, Any, Union, Tuple
from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
from airflow.models import TaskReschedule
from airflow.sensors.base import BaseSensorOperator
from airflow.utils.decorators import apply_defaults


class RichSqlSensor(BaseSensorOperator):
    """
    本 Sensor 对 airflow.operators.sensors.SqlSensor 进行了扩展
    使用方式类似原生 SqlSensor，额外支持以下特性:
        1. 支持 `pre_sql`、`post_sql` 及 `sql_on_kill` 参数，原 `sql` 参数修改为 `poke_sql`，如：

            - 在 `pre_sql` 中提交任务
                LOAD LABEL db.broker_load_label (
                    DATA INFILE ('hdfs:///file_path')
                    INTO TABLE doris_table
                    FORMAT AS 'PARQUET'
                )
                WITH BROKER 'Broker'
                PROPERTIES ('timeout': '36000');

            - 在 `poke_sql` 中轮询检查任务状态
                SHOW LOAD FROM db WHERE label = 'broker_load_label' ORDER BY CreateTime DESC LIMIT 1;

            - 在 `sql_on_kill` 中取消任务
                CANCEL LOAD FROM db WHERE LABEL = 'broker_load_label'

        2. SqlSensor 会将 *第一行第一列* 传递给 success 和 failure 函数，而本 Sensor 则会传递 poke_sql 结果集的 *第一行*，如：
            - `poke_sql` 查询 Broker load 状态后，判断 Status(第 3 个单元格) 是否为 'FINISHED'，如果是则任务失败，否则继续轮询
            success=lambda r: r[2] == 'FINISHED'
            - 个别 Broker Load 任务可能当天分区没有任何数据需要导入，Doris 会报错 'all partitions have no load data'，此时可以将这种情况置为任务成功
            success=lambda r: (r[2] == 'FINISHED' or "all partitions have no load data" in r[7]),

        3. failure 函数允许返回 2 个值，其中第 2 个值为错误信息(str)，默认行为是将第一列转化为 Bool 判断为 False
            `poke_sql` 查询 Broker load 状态后，判断 Status(第 3 个单元格) 是否为 'CANCELLED'，如果是则任务失败，并将 ErrorMsg(第 8 个单元格) 作为错误信息抛出
            failure=lambda r: (r[2] is not None and r[2] == 'CANCELLED', str(r[7]))
    """
    template_fields = ('pre_sql', 'poke_sql', 'post_sql', 'sql_on_kill', 'conn_id',)
    template_ext = ('.sql',)

    @apply_defaults
    def __init__(self,
                 conn_id: str,
                 poke_sql: str,
                 pre_sql: str = None,
                 post_sql: str = None,
                 sql_on_kill: str = None,
                 parameters: str = None,
                 success: Callable[[Any], bool] = lambda r: bool(r[0]),
                 failure: Callable[[Any], Union[bool, Tuple[bool, str]]] = lambda r: (not bool(r[0]), f'Failed by first cell: {r[0]}'),
                 fail_on_empty: bool = False,
                 allow_null: bool = True,
                 *args, **kwargs):
        self.conn_id = conn_id
        self.pre_sql = pre_sql
        self.poke_sql = poke_sql
        self.post_sql = post_sql
        self.sql_on_kill = sql_on_kill
        self.parameters = parameters
        self.success = success
        self.failure = failure
        self.fail_on_empty = fail_on_empty
        self.allow_null = allow_null
        super(RichSqlSensor, self).__init__(*args, **kwargs)

    def _get_hook(self):
        conn = BaseHook.get_connection(self.conn_id)
        self.extra = conn.extra_dejson

        allowed_conn_type = {'google_cloud_platform', 'jdbc', 'mssql', 'mysql', 'oracle', 'postgres', 'presto', 'snowflake', 'sqlite', 'vertica'}
        if conn.conn_type not in allowed_conn_type:
            raise AirflowException(f"The connection type is not supported by SqlSensor. Supported connection types: {list(allowed_conn_type)}")
        return conn.get_hook()

    def pre_execute(self, context):
        if self.reschedule and len(TaskReschedule.find_for_task_instance(context['ti'])) > 0:
            # Reschedule 模式下不执行 Pre SQL
            self.log.info('Skip Pre SQL when reschedule.')
            return

        if not isinstance(self.pre_sql, str):
            self.log.info('No pre_sql, skip ...')
            return

        self.log.info('Pre SQL Executing: \n%s', self.pre_sql)
        hook = self._get_hook()
        first_row = hook.get_first(self.pre_sql)
        self.log.info('Pre SQL first row: \n%s', first_row)

    def post_execute(self, context, result=None):
        if self.reschedule and len(TaskReschedule.find_for_task_instance(context['ti'])) > 0:
            # Reschedule 模式下不执行 Post SQL
            self.log.info('Skip Post SQL when reschedule.')
            return

        if not isinstance(self.post_sql, str):
            self.log.info('No post_sql, skip ...')
            return

        self.log.info('Post SQL Executing: \n%s', self.post_sql)
        hook = self._get_hook()
        first_row = hook.get_first(self.post_sql)
        self.log.info('Post SQL first row: \n%s', first_row)

    def on_kill(self):
        if isinstance(self.sql_on_kill, str):
            self.log.info('On Kill SQL Executing: \n%s', self.sql_on_kill)
            hook = self._get_hook()
            try:
                first_row = hook.get_first(self.sql_on_kill)
                self.log.info('On Kill SQL first row: \n%s', first_row)
            except Exception as e:
                self.log.error(e)
        else:
            self.log.info('No sql_on_kill, skip ...')

    def poke(self, context):
        hook = self._get_hook()

        self.log.info('Poking: \n%s\n(with parameters %s)', self.poke_sql, self.parameters)
        first_row = hook.get_first(self.poke_sql, self.parameters)
        self.log.info('Poke first row: %s', first_row)
        if not first_row:
            if self.fail_on_empty:
                raise AirflowException("No rows returned, raising as per fail_on_empty flag")
            else:
                return False
        self.log.info(first_row)
        if self.failure is not None:
            if callable(self.failure):
                fail, msg = self.failure(first_row)
                if fail:
                    raise AirflowException(f"Failure criteria met. self.failure({first_row}) returned True:\n{msg}")
            else:
                raise AirflowException(f"self.failure is present, but not callable -> {self.success}")
        if self.success is not None:
            if callable(self.success):
                return self.success(first_row)
            else:
                raise AirflowException(f"self.success is present, but not callable -> {self.success}")
        if self.allow_null:
            return first_row is None or bool(first_row[0])
        return bool(first_row[0])
